from qdrant_client import QdrantClient
import config
from qdrant_client.http import models


class Qdrant:
    def __init__(self):
        super(Qdrant, self).__init__()

        self.configs = config.ConfigParser()

        # specify qdrant client
        self.client = QdrantClient(
            host=self.configs.get(key='qdrant')['host'],
            port=self.configs.get(key='qdrant')['port']
        )

    def exists_collection(self, collection_name):
        all_collections = self.client.get_collections()
        return collection_name in [c_desc.name for c_desc in all_collections.collections]

    def create_collection(self, collection_name, dim):
        if not self.exists_collection(collection_name):
            # 1. 指定矢量维度, 距离计算方式
            vectors_config = models.VectorParams(size=dim, distance=models.Distance.COSINE)

            # 2. 创建集合, 创建完自动挂载
            collection = self.client.create_collection(collection_name, vectors_config)

    # ----- document -----

    def upsert_document(self, collection_name, entities):
        # 判断该集合是否存在
        if not self.exists_collection(collection_name):
            return False

        # 将实体插入集合中
        response = self.client.upsert(collection_name, entities)

        # COMPLETED 的可以立即使用, ACKNOWLEDGED 正在队列等待写, 可能需要等一段时间
        return response.status == models.UpdateStatus.COMPLETED

    def search_document(self, collection_name, data, top_k):

        # 判断该集合是否存在 秦始皇是谁
        if not self.exists_collection(collection_name):
            return False

        print(data)

        # 发送搜索请求并获取响应
        score_points = self.client.search(
            collection_name,
            query_vector=data,
            # search_params=models.SearchParams(hnsw=128, exact=False),
            limit=top_k
        )

        document_list = []
        for point in score_points:

            document = {
                "id": point.id,
                "distance": point.score,
            }

            # point.payload 第一层字段扁平到 document
            entity = point.payload
            for field, value in entity.items():
                document[field] = value
            document_list.append(document)

        return document_list
